# -*- coding: utf-8 -*-

import cv2
import numpy as np
import os


def read_label_file(label_file_path):
    with open(label_file_path) as label_file:
        records = [[int(e) for e in line.strip().split(" ")] for line in label_file.readlines()]
    return records


def show_records(image, records):
    height = image.shape[0]
    image = image[height//2:, :, :]
    for record in records:
        assert len(record) % 2 == 0
        for i in range(0, len(record), 2):
            pt1 = record[i], record[i+1]
            pt2 = record[(i+2) % len(record)], record[(i+3) % len(record)]
            cv2.line(image, pt1, pt2, (255, 255, 0), 2)
    cv2.imshow("record", image)
    cv2.waitKey()


def generate_mask_image(records, mask_image):
    tmp = np.zeros(shape=(mask_image.shape[0]+2, mask_image.shape[1]+2),
                   dtype=np.uint8)
    for record in records:
        assert len(record) % 2 == 0
        center = [0, 0]
        for i in range(0, len(record), 2):
            center[0] += record[i]
            center[1] += record[i+1]
            pt1 = record[i], record[i+1]
            pt2 = record[(i+2) % len(record)], record[(i+3) % len(record)]
            cv2.line(mask_image, pt1, pt2, (255, 255, 255), 1)
        center[0] = int(center[0] / (len(record) / 2))
        center[1] = int(center[1] / (len(record) / 2))
        cv2.floodFill(mask_image, tmp, tuple(center), (255, 255, 255))


def save_src_and_mask_image(src_image, mask_image, folder, index):
    width = src_image.shape[1] // 2
    left_src_image, right_src_image = src_image[:, 0:width, :], src_image[:, width:, :]
    left_mask_image, right_mask_image = mask_image[:, 0:width], mask_image[:, width:]
    save_path = "{}/{:>07d}.jpg".format(folder, index)
    cv2.imwrite(save_path, left_src_image)
    save_path = "{}/{:>07d}_mask.jpg".format(folder, index)
    cv2.imwrite(save_path, left_mask_image)
    save_path = "{}/{:>07d}.jpg".format(folder, index+1)
    cv2.imwrite(save_path, right_src_image)
    save_path = "{}/{:>07d}_mask.jpg".format(folder, index+1)
    cv2.imwrite(save_path, right_mask_image)


def flip_image(src_folder, dst_folder, flip_flag):
    def get_index(path):
        start = path.rfind("/")
        end = path.rfind(".")
        return int(path[start+1: end])
    src_images_path = [os.path.join(src_folder, name) for name in os.listdir(src_folder)
                       if "jpg" in name and "mask" not in name]
    src_mask_images_path = [path.replace(".jpg", "_mask.jpg") for path in src_images_path]
    assert len(src_images_path) == len(src_mask_images_path)
    for src_image_path, src_mask_image_path in zip(src_images_path, src_mask_images_path):
        image = cv2.imread(src_image_path)
        mask_image = cv2.imread(src_mask_image_path)
        dst_image = cv2.flip(image, flip_flag)
        dst_mask_image = cv2.flip(mask_image, flip_flag)
        index = get_index(src_image_path)
        dst_image_path = "{}/{:>07d}.jpg".format(dst_folder, index)
        cv2.imwrite(dst_image_path, dst_image)
        dst_mask_image_path = "{}/{:>07d}_mask.jpg".format(dst_folder, index)
        cv2.imwrite(dst_mask_image_path, dst_mask_image)
        break


def resize_image(src_folder, dst_folder):
    def get_index(path):
        start = path.rfind("/")
        end = path.rfind(".")
        return int(path[start + 1: end])

    src_images_path = [os.path.join(src_folder, name) for name in os.listdir(src_folder)
                       if "jpg" in name and "mask" not in name]
    src_mask_images_path = [path.replace(".jpg", "_mask.jpg") for path in src_images_path]
    assert len(src_images_path) == len(src_mask_images_path)
    for src_image_path, src_mask_image_path in zip(src_images_path, src_mask_images_path):
        src_image = cv2.imread(src_image_path)
        dst_image = cv2.resize(src_image, (320, 180))
        mask_image = cv2.imread(src_mask_image_path)
        dst_mask_image = cv2.resize(mask_image, (320, 180))
        _, dst_mask_image = cv2.threshold(dst_mask_image, 100, 1, cv2.THRESH_BINARY)
        index = get_index(src_image_path)
        dst_image_path = "{}/{:>07d}.jpg".format(dst_folder, index)
        cv2.imwrite(dst_image_path, dst_image)
        dst_mask_image_path = "{}/{:>07d}_mask.jpg".format(dst_folder, index)
        cv2.imwrite(dst_mask_image_path, dst_mask_image)


def main():
    # src_folder = "/home/lijun/Dataset/highway"
    # images_path = [os.path.join(src_folder, name) for name in os.listdir(src_folder) if "jpg" in name]
    # dst_folder = "/home/lijun/Dataset/vehicle_segment"
    # index = 1
    # for image_path in images_path:
    #     image = cv2.imread(image_path)
    #     image = image[image.shape[0] // 2:, :, :]
    #     label_file_path = image_path.replace("jpg", "txt")
    #     records = read_label_file(label_file_path)
    #     mask_image = np.zeros(shape=image.shape[0:2], dtype=np.uint8)
    #     generate_mask_image(records, mask_image)
    #     save_src_and_mask_image(image, mask_image, dst_folder, index)
    #     index += 2

    # src_folder = "/home/lijun/Dataset/vehicle_segment/origin"
    # dst_folder = "/home/lijun/Dataset/vehicle_segment/horizontal_flip"
    # flip_image(src_folder, dst_folder, 1)

    # src_folder = "/home/lijun/Dataset/vehicle_segment/origin"
    # dst_folder = "/home/lijun/Dataset/vehicle_segment/vertical_flip"
    # flip_image(src_folder, dst_folder, -1)

    src_folder = "/home/lijun/Dataset/vehicle_segment/horizontal_flip"
    dst_folder = "/home/lijun/Dataset/vehicle_segment/tmp"
    resize_image(src_folder, dst_folder)


if __name__ == "__main__":
    main()
